from scipy import load, savez,array,isnan,ones,arange,unique,shape,nan,concatenate,nanmax,diff,copy,argmin,argmax,multiply,sum,ceil,floor,polyfit,nanargmin, roll,nanargmax,  nanargmin, linspace, nanmin, nanmax,histogram2d,log10,mean,argsort

from pylab import plot, xlabel,ylabel, imshow,gca,xlim,find,axes,twinx,legend
import pylab

import constants as c
import tools.calc
from tools.plotroutines import *
from tools.scaledata import scaledata
from FowlerNordheim import FowlerNordheim
class IV():
  def __init__(self):
    """
    Class for keeping IV data
    """
    self.I=[[nan,nan,nan,nan],[nan,nan,nan,nan]]
    self.Imean=[[nan,nan,nan,nan],[nan,nan,nan,nan]]
    self.Ihist = [[nan,nan,nan,nan],[nan,nan,nan,nan]]
    self.Istd=[[nan,nan,nan,nan],[nan,nan,nan,nan]]
    self.VmuxOffset  = [0,0]
  def rotate(self,nshift=0,mux=1):
    """
    Rotates I and Imean of mux nshift
    """
    try:
      self.I[mux][0]=roll(self.I[mux][0],nshift,axis=2)
      self.Imean[mux][0]=roll(self.Imean[mux][0],nshift,axis=1)
    except IndexError:
      print "No data for mux:%i" % mux
    
  def zerobiasconductance(self,dV=0.02,mux=0,dI=0):
    """
    fits conductance between +- dV, save into self.Gfit
    """
    fitrange = find((self.V>-dV)*(self.V<dV))
    fit = polyfit(self.V[fitrange],self.Imean[mux][dI][:,fitrange].T,1)
    self.Gfit = fit[0,:]/77.e-6
    self.zerobias = fit[0,:]
    return self.zerobias
    
  def Icut(self,icut1 = 1,icut2 = 1):
    """
    remove data from serie which have I higher than icut (1,2 for multiplexer 1,2)
    """
    print "Icut value = %.2e" % icut1
    self.Imean[0][0][abs(self.Imean[0][0])>icut1] = nan
    self.Imean[1][0][isnan(self.Imean[0][0])*(abs(self.Imean[1][0])>icut1)] = nan
    
  def correct_offset(self,dV=0.02):
    """
    Correct offset of current and voltage device by finding the intersect of IV curves with highes and smallest slope.
    """
    I_Vzero = mean(self.Imean[0][0][:,self.V==self.V[nanargmin(abs(self.V))]],1)
    self.zerobiasconductance(dV)
    amax=argsort(abs(self.Gfit))[-20:]
    amin=argsort(abs(self.Gfit))[0]
    I_offset = I_Vzero[amin]
    tmp= (I_Vzero[amin] - I_Vzero[amax])/(self.zerobias[amax])
    #V_offset=nansum(tmp)/sum(isnan(tmp)==False)
    V_offset = tools.calc.nanmean(tmp)
    print "$V_{offset}=%.1f mV, I_{offset}=%.3e A" % (V_offset*1e3,I_offset)
    if not isnan(V_offset):
      self.V = self.V-V_offset #self.data[:,:,0] = self.data[:,:,0]-V_offset
      self.Imean[0][0]=self.Imean[0][0]-I_offset
      self.I[0][0]=self.I[0][0]-I_offset
      
  def histI(self,mux=0,length=300,conf=.75):
    """
    Calculate histograms of I data with shape aproximately shape length x length
    """
    V=self.V+self.VmuxOffset[mux]
    dV=(V[1]-V[0])
    self.binV = arange(V.min(),V.max(),abs(dV*(1+(nanargmin(V)-nanargmax(V))//length)))
    X = V[self.Ns[0]]
    self.binI = [nan,nan]
    for dI in [0,1]:
      Cmin, Cmax = getclim(self.Imean[mux][dI], conf=conf)
      self.binI[dI]  = linspace(Cmin,Cmax,length)
      self.Ihist[mux][dI] = array([histogram2d(X,Y,bins=(self.binV,self.binI[dI]))[0] for Y in self.Imean[mux][dI][:,self.Ns[0]]])
    
  def plot2Dgroup(self,N,mux=0,dI=1,**kw):
    """
    Plot  results from histI
    """
    im = self.Ihist[mux][dI][N].sum(0)
    imshow(im.T if sum(im)==0 else log10(im.T),origin='lower',extent=(nanmin(self.binV),nanmax(self.binV),nanmin(self.binI[dI]),nanmax(self.binI[dI])),aspect='auto',**kw)
    xlabel("Bias [V]");ylabel("dI/dV [S]" if dI==1 else "I [A]")
    
  def groupIV(self,dV1=0.02,dV2=0.2,lims=[.9,1.1]):
    """
    Groups the IVs into tree groups (IV.Group):
    
    IV.ratio > 1 : Zero bias enhancement
    IV.ratio ~ 1 : Aprx. Linear
    IV.ratio < 1 : gap
    """
    self.ratio = self.zerobiasconductance(dV1)/self.zerobiasconductance(dV2)
    self.Group = {}
    self.Group[c.ZEROBIAS_ENHANCE]   = find(self.ratio>=lims[1])
    self.Group[c.LINEAR]             = find((self.ratio>lims[0]) *(self.ratio<lims[1]))
    self.Group[c.GAP]                = find(self.ratio<=lims[0])
    
  def stdIVs(self,STDmax=0.05,mux=0,dI=0):
    """
    1) Calculate the standard derivative of I (IV.Istd)
    2) Finds IVs with low STD normalized to max I (IV.Stable)
    """
    try:
      self.Istd[mux][dI] = self.I[mux][dI].std(1)
    except NameError,AttributeError:
      self.Istd[mux][dI] = self.I[mux][dI].std(1)
    self.Stable = tools.calc.nanmean(self.Istd[mux][dI],axis=1)/nanmax(self.Imean[mux][dI],axis=1)
  def __numDI__(self,nfitlength=9,mux=0,dI=1):
    self.Imean[mux][dI] = tools.calc.slope_diff(self.V,self.I[mux][0],nfitlength=nfitlength,dI=dI)
  def numDI(self,nfitlength=9,mux=0,**kw):
    self.nfitlength = nfitlength
    self.__numDI__(nfitlength=nfitlength,mux=mux,dI=1)
    if mux==0:
      self.__numDI__(nfitlength=nfitlength*3,mux=mux,dI=2)
    if len(kw)>0:
      print "These arguments to numDI are old and are not used anymore:"
      print kw
  def FowlerNordheim(self,**kw):
    """
    Analyse IV using Fowler Nordheim method. Finds minimums :
    """
    self.FN = FowlerNordheim(self.V,self.Imean[0][0],**kw)
  def savez(self,filename=None):
    """ 
    Save object in npz file.
    """
    if vars(self).has_key('filename') and filename is None:
      filename = self.filename
    else:
      self.filename=filename
    if filename is None:
      print """ERROR
             No filename given"""
      return 0
    savez(filename,vars(self))
    #savez(filename,data=self.data,motor=self.motor,meta=self.meta,I=self.I,V=self.V,Ns=self.Ns,SavedTraces=self.SavedTraces)
    
  def loadz(self,filename=None):
    """ 
    Loads object (and data) from saved npz file.
    """
    if vars(self).has_key('filename') and filename is None:
      filename = self.filename
    if filename is None:
      print """ERROR
             No filename given"""
      return 0
    tmp=load(filename)
    for i,j in tmp['arr_0'].tolist().iteritems():
      setattr(self,i,j)
   
  def PlotIV(self,dat_i,mux=0,dI=1,ns=None,zeroalign=False,normalize=0,mV=True,**kw):
    """
    Plot IV(s) with number(s) dat_i
    
    """
    Y =  - (self.Imean[mux][dI][dat_i].T[argmin(abs(self.V)),:] if zeroalign else 0)
    X = (1.e3 if mV else 1.)*(self.V+self.VmuxOffset[mux])
    if ns is None:
      Y += self.Imean[mux][dI][dat_i].T
      if normalize<>0:
        Y = Y/nanmax(Y*normalize,0)
    else:
      Y += self.Imean[mux][dI][dat_i][:,self.Ns[ns]].T
      X = X[self.Ns[ns]]
      if normalize<>0:
        Y = Y/nanmax(Y*normalize,0)
    plot(X,Y,**kw)
    ylabel('I [A]' if mux+dI==0 else ('dI/dV [S]' if mux+dI==1 else '$d^2I/dV^2 [S/V]$' ))
    xlabel('mV'if mV else 'V')
    mysciscale()
    
  def PlotIM(self,mux=0,dI=0,nb=0,av=False,maxV=None,traces=None,normalize=True,reorderV=True,single=False,nVg=1,nVb=5,log=False,circular=False,**kw):
    """
    Plot image of data from 'mux' the 'dI' derivative
    mux = 0 (default),1
    dI = 0 (default),1,2
    nb = 0 (default),1
    reorderV = True (default) / False , True: split the V axis into two independent IV traces, False: keep the originate seqeunce of the IV
    single = True / False  (default) , True Plot only a single trace, use 'nb'.
    nVg = integer (default 1): roll avarage the 'gate' with a window of 'nVg'
    nVb = integer (default 5): roll avarage the 'bias' with a window of 'nVb'
    circular = False (default)/ True: Do the rolling avarage in a circular manor
    log = False (default)/ True: plot the logaritm (base 10) of the data.
    traces = list / None(default): if None plot all IV traces otherwize only plot traces in list "traces"
    normalize = True (default) / False , If True normalize each trace to the value at V=0 if dI+mux==1 else to maximum of trace
    """
    im = self.Imean[mux][dI]
    im = tools.calc.rollavg(im,n=nVb,circular=circular,axis=1)
    im = tools.calc.rollavg(im,n=nVg,circular=circular)
    if single:
      if av:
        im = nansum([im[:,ns] for ns in self.Ns],axis=0)/nansum([~isnan(im[:,ns]) for ns in self.Ns],axis=0)
      else:
        im = im[:,self.Ns[nb]]
      V = self.V[self.Ns[0]]
    elif reorderV:
      N=array([nan]*len(im)).reshape((len(im),1))
      im = concatenate((im[:,self.Ns[0]],N,im[:,self.Ns[1]]),axis=1)
      V = concatenate((self.V[self.Ns[0]],[nan],self.V[self.Ns[1]]))
    else:
      V = self.V
    V += self.VmuxOffset[mux]
    if traces is not None:
      im = im[traces,:]
    if maxV is not None:
      tmp = find(abs(V)<maxV)
      im = im[:,tmp]
      V = V[tmp]
    Gfit = im[:,nanargmin(abs(V))]
    if normalize:
      norm = 1./Gfit if (mux+dI)==1 else 1./nanmax(im,1)
      im = multiply(norm,im.T).T
    Xlim = (min(traces)-.5,max(traces)+.5) if traces is not None else (-.5,len(im)+.5)
    kw['interpolation'] =  kw['interpolation'] if kw.has_key ('interpolation') else 'nearest'
    if sum(~isnan(im))>0:
      imshow(log10(abs(im.T)) if log else im.T,origin='lower',extent=(Xlim[0],Xlim[1],0,4),aspect='auto',**kw)#extent=(0,len(im),)
    else:
      print 'only NaN\'s'
    ylabel('$V_{bias} / mV$')
    test = lambda x,n=8.: [1,2,5,25,50,100,200,250,500,1000][find(array([x//j for j in [1,2,5,25,50,100,200,250,500,1000]])<n)[0]]
    V *= 1.e3
    ax1 = gca()
    ax1.yaxis.set_ticks(range(5))
    n=8. if single else 4.
    DeltaV = test((max(V)-min(V)),n)
    ticvals = arange(ceil((min(V)-DeltaV/6.)/DeltaV)*DeltaV,max(V)+DeltaV/6.,DeltaV)
    tmp = array([V[nanargmin(abs(V-v))]==V for v in ticvals]).any(0)
    ticvals = (V)[tmp].round()
    mytics = find(tmp)*4./len(V)
    n=find(abs(diff(roll(mytics,1)))<.2)
    mytics,ticvals = mytics.tolist(),ticvals.tolist()
    for i in n:
      del(mytics[i])
      del(ticvals[i]) 
    ax1.yaxis.set_ticks(mytics)
    ax1.yaxis.set_ticklabels(ticvals)
    pylab.ylim(0,4)
    autoclim()
    scaledata(Yarray=V)
    xlim(Xlim)
 
  def findgapspacing(self,peaks,filterpeak=4,Vtol=20e-3,mux=0,dI=1):
    gaps = []
    self.Allgaps = []
    for i in range(len(self.Imean[mux][dI])):
      dI0 = self.Imean[mux][dI][i][argmin(abs(self.V))]
      p = array(filter(lambda p: p[3]>filterpeak*dI0, peaks[find(peaks[:,0]==i),:]))
      if len(p)==0:
        continue
      elif sum(p[:,2]==0)==0 or sum(p[:,2]==1)==0:
        continue
      tmp=[nan,nan]
      for j in [0,1]:
        ptmp = p[p[:,2]==j,1]
        p_pos = p[p[:,2]==j][ptmp>0]
        tmp[j] = [nan,nan] if len(p_pos)==0 else [p_pos[argmin(p_pos[:,1]),1],p_pos[argmin(p_pos[:,1]),3]-dI0]
        p_neg = p[p[:,2]==j][ptmp<0]
        tmp[j] += [nan,nan] if len(p_neg)==0 else [p_neg[argmax(p_neg[:,1]),1],p_neg[argmax(p_neg[:,1]),3]-dI0]
        tmp[j] += [dI0,i,j]
        self.Allgaps.append(tmp)
      if (abs(tmp[0][0]-tmp[1][0])<Vtol) and (abs(tmp[0][2]-tmp[1][2])<Vtol):
        gaps.append(tmp[0])
        gaps.append(tmp[1])
    gaps = array(gaps)
    return gaps
    
  def PlotSingle(self,nr,rotate_n=0,nfitlength=9,twin=True,muxplot=[0]):
    """
    Plot a single IV with number 'nr'. Used to find parameters for entire serie. data in mux==1 is rotated with 'rotate_n', and the slope are calculated using  a window of 'nfitlength'
    muxplot: plot mux in list muxplot (default [0])
    """
    lines=[]
    ax1 = axes([0.1,0.1,0.8,0.8]) if twin else axes([0.1,0.7,0.8,0.3])
    ax1.xaxis.set_visible(False)
    ax1.yaxis.set_ticklabels([""]*len(ax1.yaxis.get_ticklabels()))    
    ax2 = twinx() if twin else  axes([0.1,0.4,0.8,0.3],sharex=ax1)
    ax2.xaxis.set_visible(False)
    ax2.yaxis.set_ticklabels([""]*len(ax2.yaxis.get_ticklabels()))
    ax3=twinx() if twin else axes([0.1,0.1,0.8,0.3],sharex=ax1)
    ax3.xaxis.set_visible(True)
    ax3.yaxis.set_ticklabels([""]*len(ax3.yaxis.get_ticklabels()))
    I0= self.I[0][0][[nr]]
    Id0 = tools.calc.slope_diff(self.V,I0,nfitlength=nfitlength,dI=1)[0]
    if muxplot.count(0)==1:
      lines.append((ax1.plot(self.V,I0.mean(1)[0],'-r',label='MUX1'),'MUX1'))
      lines.append((ax3.plot(self.V,tools.calc.slope_diff(self.V,I0,nfitlength=nfitlength*3,dI=2)[0],':k',lw=2,label='diff2 MUX1'),'diff2 MUX1'))
    if muxplot.count(1)==1:
      I1= roll(self.I[1][0][[nr]],rotate_n,axis=2)
      Id1 = tools.calc.slope_diff(self.V,I1,nfitlength=nfitlength,dI=1)[0]
      lines.append((ax2.plot(self.V,I1.mean(1)[0],'--b',label='MUX2'),'MUX2'))
      lines.append((ax3.plot(self.V,Id1,'-k',lw=2,label='diff MUX2'),'diff MUX2'))
      
    if twin or not muxplot.count(1)==1:
      lines.append((ax2.plot(self.V,Id0,'-.g',label='diff MUX1'),'diff MUX1'));
    if twin:
      legend (array(lines)[:,0],array(lines)[:,1],loc='best')
    ax1.set_ylabel('$I [A]$')
    ax2.set_ylabel('$dI/dV  [S]$')
    ax3.set_ylabel('$d^2I/dV^2  [A/V^2]$')
    xlabel('U [V]')